package org.sunyaxing.transflow.transflowapp.plugins.filters;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;
import com.alibaba.fastjson2.JSONObject;
import com.openai.client.OpenAIClient;
import com.openai.client.okhttp.OpenAIOkHttpClient;
import com.openai.models.chat.completions.ChatCompletion;
import com.openai.models.chat.completions.ChatCompletionCreateParams;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sunyaxing.transflow.api.TransData;
import org.sunyaxing.transflow.api.common.Handle;
import org.sunyaxing.transflow.api.common.ano.Extension;
import org.sunyaxing.transflow.api.common.ano.HandleItem;
import org.sunyaxing.transflow.api.common.ano.JobParamItem;
import org.sunyaxing.transflow.api.common.ano.ScopeContentCheck;
import org.sunyaxing.transflow.api.extensions.base.ExtensionContext;
import org.sunyaxing.transflow.api.extensions.base.typesimpl.TransFlowMiddleJsonGatewayHandler;
import org.sunyaxing.transflow.transflowapp.controllers.JobLogsWs;

import java.util.List;
import java.util.Optional;
import java.util.function.Function;

@Extension("plugin-openai-filter")
@ScopeContentCheck(
        value = {
                @JobParamItem(field = "token", label = "token"),
                @JobParamItem(field = "baseUrl", label = "baseUrl", defaultValue = "https://dashscope.aliyuncs.com/compatible-mode/v1"),
                @JobParamItem(field = "temperature", label = "temperature", defaultValue = "0.1"),
                @JobParamItem(field = "prompt", label = "prompt", type = "code", defaultValue = ""),
                @JobParamItem(field = "maxTokens", label = "maxTokens", defaultValue = "8196"),
                @JobParamItem(field = "askKey", label = "用户提问KEY", defaultValue = "userMsg"),
                @JobParamItem(field = "responseKey", label = "AI返回KEY", defaultValue = "contentStr"),
        },
        handle = @HandleItem(field = "model", label = "model", required = true, defaultValue = "")
)
public class OpenAiTransflowFilterExt extends TransFlowMiddleJsonGatewayHandler {
    private static final Logger log = LoggerFactory.getLogger(OpenAiTransflowFilterExt.class);
    private OpenAIClient client;
    private String prompt;
    private Double temperature;
    private Long maxTokens;
    private String askKey;
    private String responseKey;

    public OpenAiTransflowFilterExt(ExtensionContext extensionContext) {
        super(extensionContext);
    }

    @Override
    protected void afterInitHandler(JSONObject config, List<Handle> handles) {
        String token = config.getString("token");
        String baseUrl = config.getString("baseUrl");
        this.temperature = config.getDouble("temperature");
        this.prompt = config.getString("prompt");
        this.maxTokens = config.getLong("maxTokens");
        this.askKey = config.getString("askKey");
        this.responseKey = config.getString("responseKey");
        this.client = OpenAIOkHttpClient.builder()
                .baseUrl(baseUrl)
                .apiKey(token)
                .build();
    }

    @Override
    public Function<TransData<JSONObject>, Boolean> parseHandleToConsumer(String handleId, String model) {
        return transData -> {
            JSONObject inputMsg = transData.getData();
            ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
                    .addSystemMessage(prompt)
                    .addUserMessage(inputMsg.getString(askKey))
                    .model(model)//qwen-plus
                    .temperature(temperature)
                    .maxCompletionTokens(maxTokens)
                    .build();
            ChatCompletion chatCompletion = client.chat().completions().create(params);
            chatCompletion.choices().forEach(choice -> {
                Optional<String> content = choice.message().content();
                content.ifPresent(contentStr -> {
                    JobLogsWs.sendMessage(jobId, "【{}】AI处理结果 {}", extensionContext.getNodeName(), contentStr);
                    contentStr = HtmlUtil.removeHtmlTag(contentStr, "think");
                    if (contentStr.contains("</think>")) {
                        contentStr = CollectionUtil.getLast(StrUtil.split(contentStr, "</think>"));
                    }
                    inputMsg.put(responseKey, contentStr);
                });
            });
            return true;
        };
    }

    @Override
    public void destroy() {
        this.client.close();
    }
}
